package net

import (
	"context"
	"sync"

	"google.golang.org/grpc"
	"google.golang.org/grpc/connectivity"
	"google.golang.org/grpc/credentials/insecure"
)

var (
	connections = map[string]*grpc.ClientConn{}
	mu          = &sync.Mutex{}
)

func GetConn(address string) (*grpc.ClientConn, error) {
	mu.Lock()
	defer mu.Unlock()

	if conn, ok := connections[address]; ok {
		if conn.GetState() == connectivity.Ready {
			return conn, nil
		}
	}

	conn, err := grpc.Dial(address, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return conn, err
	}

	connections[address] = conn

	return conn, err
}
func OnConnectionOpened(ctx context.Context, conn *grpc.ClientConn, callback func(conn *grpc.ClientConn)) {
	states := []connectivity.State{connectivity.Ready}
	onConnectionStateChanged(ctx, conn, callback, states...)
}
func OnConnectionClosed(ctx context.Context, conn *grpc.ClientConn, callback func(conn *grpc.ClientConn)) {
	states := []connectivity.State{connectivity.Idle, connectivity.Shutdown, connectivity.TransientFailure}
	onConnectionStateChanged(ctx, conn, callback, states...)
}
func onConnectionStateChanged(ctx context.Context, conn *grpc.ClientConn, callback func(conn *grpc.ClientConn), states ...connectivity.State) {
	done := make(chan struct{})

	go func() {
		defer close(done)

		for {
			state := conn.GetState()
			change := conn.WaitForStateChange(ctx, state)
			if !change {
				return
			}

			currentState := conn.GetState()
			for _, s := range states {
				if currentState == s {
					callback(conn)
					return
				}
			}
		}
	}()
}
